from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_graphviz
import graphviz


def main():
    wine = load_wine()
    x_train, x_test, y_train, y_test = train_test_split(wine.data, wine.target)
    clf = DecisionTreeClassifier(criterion="gini", splitter="best")  # criterion="entropy" 标准=信息熵
    clf.fit(x_train, y_train)  # 训练
    train_score = clf.score(x_train, y_train)
    test_score = clf.score(x_test, y_test)
    print("train score:", train_score)
    print(" test score:", test_score)
    wine_decision_tree = export_graphviz(clf, out_file=None,  # 输出文件的路径，如果为None，则返回图形描述字符串。
                                         feature_names=wine.feature_names,
                                         class_names=wine.target_names,
                                         filled=True, rounded=True,  # 表示用不同颜色填充节点以表示不同的类别，使节点的边框为圆形。
                                         special_characters=True)  # 允许在节点标签中使用特殊字符。
    graph = graphviz.Source(wine_decision_tree)  # 将图形描述转换为图形对象
    graph.render("wine_decision_tree", view=True)  # render方法将图形保存为文件，view=True自动打开查看生成的图形文件。


if __name__ == '__main__':
    main()
